import argparse
import os
import subprocess
import whisper
from datetime import datetime
from tempfile import NamedTemporaryFile


def extract_audio(video_path: str) -> str:
    temp_audio = NamedTemporaryFile(suffix=".mp3", delete=False)
    try:
        subprocess.run(["ffmpeg", "-i", video_path, "-b:a", "192k", "-q:a", "0", "-map", "a", temp_audio.name], check=True)
        return temp_audio.name
    except subprocess.CalledProcessError as e:
        os.remove(temp_audio.name)
        raise RuntimeError(f"音频提取失败: {str(e)}")


def transcribe_to_srt(audio_path: str, model_size: str, target_language: str) -> str:
    try:
        model = whisper.load_model(model_size)
        result = model.transcribe(audio_path, task='translate' if target_language != 'en' else 'transcribe')

        with NamedTemporaryFile(suffix=".srt", delete=False) as temp_file:
            srt_path = temp_file.name

            with open(srt_path, 'w', encoding='utf-8') as f:
                for i, segment in enumerate(result['segments'], 1):
                    start = datetime.fromtimestamp(segment['start']).strftime('%H:%M:%S,%f')[:-3]
                    end = datetime.fromtimestamp(segment['end']).strftime('%H:%M:%S,%f')[:-3]
                    f.write(f"{i}\n{start} --> {end}\n{segment['text'].strip()}\n\n")

            return srt_path

    except Exception as e:
        os.remove(audio_path)
        raise RuntimeError(f"语音转写失败: {str(e)}")


def main():
    parser = argparse.ArgumentParser(description='将视频转换为SRT字幕文件')
    parser.add_argument('video_path', type=str, help='视频文件路径')
    parser.add_argument('--model_size', type=str, default='base',
                        help='Whisper模型规模 (tiny, base, small, medium, large)')
    parser.add_argument('--target_language', type=str, default='zh',
                        help='目标语言代码 (如zh, en, ja等)')
    args = parser.parse_args()

    try:
        print("正在提取音频...")
        audio_path = extract_audio(args.video_path)
        print("音频提取成功，正在转写字幕...")
        srt_path = transcribe_to_srt(audio_path, args.model_size, args.target_language)
        print(f"字幕转写成功，文件已保存至: {srt_path}")
    except Exception as e:
        print(f"错误: {str(e)}")


if __name__ == '__main__':
    main()